# -*- coding: utf-8 -*-
import networkx as nx
# import matplotlib.pylab as plt

'''
拓扑计算，生成
'''


class PathGenerator:
    def __init__(self):
        self.G = nx.Graph()

    def add_edge(self, u, v, weight=None):
        '''
        向图中添加一条边
        :param u:
        :param v:
        :param weight: 权重
        :return:
        '''
        try:
            self.G.add_edge(u, v, weight=weight)
        except:
            return False
        return True

    def check_edge(self, u, v):
        '''
        检查边是否存在
        '''
        return self.G.has_edge(u, v)

    def remove_edge(self, u, v):
        '''
        删除一条边
        :param u:
        :param v:
        :return:
        '''
        try:
            self.G.remove_edge(u, v)
        except:
            return False
        return True

    def generate_shortest_path(self, source, target=None):
        '''
        计算指定源节点到其他节点的最短路,如未指定目的节点，则返回到其他全部节点的路径
        :param source: 源节点
        :return: 
        如果指定了目的节点 则返回路径列表
        如果未指定目的节点，则返回以目的节点为键，路径为值的字典
        '''
        try:
            return nx.shortest_path(self.G, source=source, target=target)
        except:
            return None

    def generate_all_shortest_path(self):
        '''
        计算全拓扑节点两两间最短路
        :param source: 源节点
        :return: 字典的字典 第一层字典键为源节点 第二层字典的键为目的节点，值为路径
        '''
        try:
            return dict(nx.all_pairs_shortest_path(self.G))
        except:
            return None

    def get_nodes(self):
        '''
        获取全部节点
        '''
        return [n for n in self.G.nodes]

    def remove_node(self, n):
        '''
        删除指定节点
        :param n: 要删除的节点名
        '''
        try:
            self.G.remove_node(n)
            return True
        except:
            return False

    def clean_topo(self):
        '''
        清空拓扑
        '''
        self.G.clear()

    # def draw_map(self):
    #     '''
    #     绘图，调试用
    #     '''
    #     nx.draw(self.G, with_labels=True)
    #     plt.show()


if __name__ == '__main__':
    pg = PathGenerator()
    pg.add_edge(1, 2)
    pg.add_edge(2, 3)
    pg.add_edge(3, 4)
    pg.add_edge(1, 5)
    pg.add_edge(3, 5)
    print(pg.generate_shortest_path(3))
    pg.remove_edge(3, 5)
    print(pg.generate_shortest_path(3))
    print(pg.generate_shortest_path(3, 5))
    print(pg.generate_all_shortest_path())
